# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

from absl import app, flags

import pandas
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

FLAGS = flags.FLAGS


def plot_book_trade(book_df, trade_df, plot_out_root):
  plt.figure(figsize=(20, 10))

  book_index = pandas.DatetimeIndex(book_df['timestamp'])
  trade_index = pandas.DatetimeIndex(trade_df['timestamp'])

  plt.plot(trade_index,
           trade_df['price'],
           'kX',
           book_index,
           book_df['bid0'],
           'ro-',
           book_index,
           book_df['ask0'],
           'bo-',
           lw=0.2,
           markersize=1)
  plt.title('20181107 BTC-JPY.IMMEDIATE.Bitflyer.BTC_JPY')
  plt.savefig('%s/book_trade_plot' % plot_out_root)
  plt.close()


def extract_part(book_df, trade_df):
  start_ts = book_df.iloc[0]['timestamp']
  end_ts = start_ts + 2 * 60 * 1e9  # 5 min
  book_df = book_df.loc[book_df['timestamp'] <= end_ts]
  trade_df = trade_df.loc[(trade_df['timestamp'] <= end_ts) & (trade_df['timestamp'] >= start_ts)]
  return book_df, trade_df


def read_csv_into_df(csv_path):
  df = pandas.read_csv(csv_path, sep=',', header=0)
  df['datetime'] = pandas.to_datetime(df['timestamp'], unit='ns')
  df = df.set_index('datetime')
  return df


def main(argv):
  book_csv_path = FLAGS.book_csv_path
  assert book_csv_path, '--book_csv_path must be specified.'

  trade_csv_path = FLAGS.trade_csv_path
  assert trade_csv_path, '--trade_csv_path must be specified.'

  plot_out_root = FLAGS.plot_out_root
  assert plot_out_root, '--plot_out_root must be specified.'

  book_df = read_csv_into_df(book_csv_path)
  trade_df = read_csv_into_df(trade_csv_path)

  # book_df, trade_df = extract_part(book_df, trade_df)
  plot_book_trade(book_df, trade_df, plot_out_root)

  return 0


if __name__ == '__main__':
  flags.DEFINE_string('book_csv_path', None, 'Input csv file path.')

  flags.DEFINE_string('trade_csv_path', None, 'Input csv file path.')

  flags.DEFINE_string('plot_out_root', None, 'Output plot file root path.')
  app.run(main)
